from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import json
import asyncio
import numpy as np
from typing import List, Dict, Optional
import os
from dotenv import load_dotenv

from llm_service import LLMService
from recommendation import RecommendationEngine

load_dotenv()

app = FastAPI(title="Knowledge Graph Recommendation System")

# CORS配置
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 初始化服务
llm_service = LLMService(
    api_key=os.getenv("QWEN_API"),
    base_url=os.getenv("QWEN_API_BASE_URL"),
    model=os.getenv("QWEN_MODEL")
)
recommendation_engine = RecommendationEngine()

# 加载演示数据
with open("../demo_data.json", "r", encoding="utf-8") as f:
    demo_data = json.load(f)


class PostRequest(BaseModel):
    content: str
    user_id: str


class WebSocketManager:
    def __init__(self):
        self.active_connections: List[WebSocket] = []

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)

    def disconnect(self, websocket: WebSocket):
        self.active_connections.remove(websocket)

    async def broadcast(self, message: dict):
        for connection in self.active_connections:
            try:
                await connection.send_json(message)
            except:
                pass


ws_manager = WebSocketManager()


async def find_similar_topics_by_embedding(topic: str, nodes: List[Dict], llm_service: LLMService, top_k: int = 3) -> List[Dict]:
    """使用embedding向量找到相似的话题节点

    返回: [{"node_id": str, "label": str, "similarity": float}, ...]
    """
    try:
        # 获取输入话题的embedding
        topic_embedding = await llm_service.get_embedding(topic)

        similarities = []

        for node in nodes:
            if node["type"] != "topic":
                continue

            # 获取节点标签的embedding
            node_embedding = await llm_service.get_embedding(node["label"])

            # 计算余弦相似度
            similarity = llm_service.cosine_similarity(topic_embedding, node_embedding)

            if similarity > 0.6:  # 相似度阈值
                similarities.append({
                    "node_id": node["id"],
                    "label": node["label"],
                    "similarity": similarity
                })

        # 按相似度降序排序
        similarities.sort(key=lambda x: x["similarity"], reverse=True)

        return similarities[:top_k]
    except Exception as e:
        print(f"相似度计算错误: {e}")
        return []


@app.get("/")
async def root():
    return {"message": "Knowledge Graph Recommendation System API"}


@app.get("/api/data")
async def get_demo_data():
    """获取演示数据"""
    return demo_data


@app.get("/api/graph")
async def get_graph():
    """获取知识图谱数据"""
    return demo_data["knowledge_graph"]


@app.get("/api/users")
async def get_users():
    """获取用户列表"""
    return demo_data["users"]


@app.get("/api/posts")
async def get_posts():
    """获取帖子列表"""
    return demo_data["posts"]


@app.post("/api/analyze")
async def analyze_post(request: PostRequest):
    """分析帖子内容"""
    try:
        analysis = await llm_service.analyze_post(request.content)
        return {
            "success": True,
            "analysis": analysis
        }
    except Exception as e:
        return {
            "success": False,
            "error": str(e)
        }


@app.post("/api/recommend")
async def get_recommendations(request: PostRequest):
    """获取推荐"""
    try:
        # 分析帖子
        analysis = await llm_service.analyze_post(request.content)

        # 获取推荐
        recommendations = recommendation_engine.find_recommendations(
            topics=analysis.get("topics", []),
            user_id=request.user_id,
            graph_data=demo_data["knowledge_graph"],
            users=demo_data["users"],
            posts=demo_data["posts"]
        )

        return {
            "success": True,
            "analysis": analysis,
            "recommendations": recommendations
        }
    except Exception as e:
        return {
            "success": False,
            "error": str(e)
        }


@app.post("/api/node-embedding")
async def get_node_embedding(request: dict):
    """获取节点的embedding向量"""
    try:
        node_label = request.get("label")
        if not node_label:
            return {
                "success": False,
                "error": "Missing node label"
            }

        # 获取embedding
        embedding = await llm_service.get_embedding(node_label)

        return {
            "success": True,
            "label": node_label,
            "embedding": embedding[:10],  # 只返回前10维供显示
            "dimension": len(embedding),
            "norm": float(np.linalg.norm(embedding))
        }
    except Exception as e:
        return {
            "success": False,
            "error": str(e)
        }


@app.websocket("/ws/graph-stream")
async def websocket_endpoint(websocket: WebSocket):
    """WebSocket实时推送"""
    await ws_manager.connect(websocket)

    try:
        while True:
            # 接收客户端消息
            data = await websocket.receive_json()

            if data.get("type") == "post":
                post_content = data.get("content")
                user_id = data.get("user_id", "demo_user")

                # 步骤1: 发送处理开始消息
                await websocket.send_json({
                    "type": "processing_start",
                    "message": "开始分析帖子..."
                })

                # 步骤2: LLM分析帖子
                analysis = await llm_service.analyze_post(post_content)
                await asyncio.sleep(0.5)

                await websocket.send_json({
                    "type": "analysis_complete",
                    "payload": analysis
                })

                # 步骤3: 处理话题节点 - 总是创建新节点并连接到相关节点
                topics = analysis.get("topics", [])
                topic_node_ids = []  # 存储新建的话题节点ID

                for i, topic in enumerate(topics):
                    await asyncio.sleep(0.3)

                    # 总是创建新节点（不检查是否已存在）
                    import time
                    timestamp = int(time.time() * 1000)
                    new_node_id = f"topic_{topic}_{user_id}_{timestamp}_{i}"

                    await websocket.send_json({
                        "type": "node_added",
                        "payload": {
                            "node": {
                                "id": new_node_id,
                                "label": topic,
                                "name": topic,
                                "type": "topic",
                                "color": "#4ECDC4",
                                "new": True
                            }
                        }
                    })
                    topic_node_ids.append(new_node_id)

                    # 使用embedding向量找到最相似的节点并创建连接（确保连通性）
                    similar_topics = await find_similar_topics_by_embedding(
                        topic,
                        demo_data["knowledge_graph"]["nodes"],
                        llm_service,
                        top_k=3  # 找到前3个最相似的节点
                    )

                    # 确保至少连接到1个节点（保证图连通）
                    if not similar_topics and len(demo_data["knowledge_graph"]["nodes"]) > 0:
                        # 如果没有找到相似节点，连接到第一个topic节点
                        fallback_node = None
                        for node in demo_data["knowledge_graph"]["nodes"]:
                            if node["type"] == "topic":
                                fallback_node = node
                                break

                        if fallback_node:
                            similar_topics = [{
                                "node_id": fallback_node["id"],
                                "label": fallback_node["label"],
                                "similarity": 0.5
                            }]

                    # 创建连接
                    for similar in similar_topics:
                        await asyncio.sleep(0.2)
                        await websocket.send_json({
                            "type": "edge_added",
                            "payload": {
                                "edge": {
                                    "source": new_node_id,
                                    "target": similar["node_id"],
                                    "type": "相关",
                                    "similarity": similar["similarity"]
                                }
                            }
                        })

                # 步骤5: 计算推荐
                recommendations = recommendation_engine.find_recommendations(
                    topics=topics,
                    user_id=user_id,
                    graph_data=demo_data["knowledge_graph"],
                    users=demo_data["users"],
                    posts=demo_data["posts"]
                )

                # 步骤6: 发送推荐链条（普通推荐用青色）
                for i, rec in enumerate(recommendations):
                    await asyncio.sleep(0.5)
                    await websocket.send_json({
                        "type": "recommendation_path",
                        "payload": {
                            "path": rec["path"],
                            "path_ids": rec.get("path_ids", []),
                            "explanation": rec["explanation"],
                            "color": "#00D9FF",  # 普通推荐用青色
                            "type": rec.get("rec_type", "normal"),
                            "is_first": (i == 0)  # 标记第一条推荐
                        }
                    })

                # 步骤7: 完成
                await websocket.send_json({
                    "type": "processing_complete",
                    "message": "推荐生成完成！"
                })

            elif data.get("type") == "request_counter_recommendation":
                user_id = data.get("user_id", "demo_user")

                await websocket.send_json({
                    "type": "processing_start",
                    "message": "正在生成反向推荐..."
                })

                # 获取反向推荐
                counter_recs = recommendation_engine.find_counter_recommendations(
                    user_id=user_id,
                    graph_data=demo_data["knowledge_graph"],
                    users=demo_data["users"],
                    posts=demo_data["posts"]
                )

                # 发送反向推荐（使用紫色）
                for i, rec in enumerate(counter_recs):
                    await asyncio.sleep(0.5)
                    await websocket.send_json({
                        "type": "recommendation_path",
                        "payload": {
                            "path": rec["path"],
                            "path_ids": rec.get("path_ids", []),
                            "explanation": rec["explanation"],
                            "color": "#FF00FF",  # 反向推荐用紫色
                            "type": "counter",
                            "is_first": (i == 0)  # 标记第一条推荐
                        }
                    })

                await websocket.send_json({
                    "type": "processing_complete",
                    "message": "反向推荐生成完成！"
                })

    except WebSocketDisconnect:
        ws_manager.disconnect(websocket)
    except Exception as e:
        print(f"WebSocket error: {e}")
        ws_manager.disconnect(websocket)


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)
